# -*- coding: utf-8 -*-
import pandas as pd
import json
from utils.llm_utils import LLMClient

class TopicLabeler:
    def __init__(self, api_key, base_url, model_name):
        self.llm_client = LLMClient(api_key, base_url, model_name)
    
    def _parse_cluster_file(self, cluster_file):
        """
        解析聚类文件的公共方法
        """
        clusters = {}
        
        # 读取文件内容
        content = cluster_file.read().decode('utf-8')
        
        # 尝试判断文件格式
        try:
            # 尝试解析为JSON
            data = json.loads(content)
            
            # 处理VOSviewer格式的JSON
            if 'network' in data and 'items' in data['network']:
                # VOSviewer格式：按cluster字段分组
                for item in data['network']['items']:
                    cluster_id = f"cluster_{item['cluster']}"
                    term = item.get('label', item.get('id', ''))
                    
                    if cluster_id not in clusters:
                        clusters[cluster_id] = []
                    clusters[cluster_id].append(term)
                
                # 将列表转换为分号分隔的字符串
                for cluster_id in clusters:
                    clusters[cluster_id] = ';'.join(clusters[cluster_id])
            
            # 处理直接的聚类结果JSON格式 
            elif isinstance(data, dict):
                # 假设JSON格式为 {"cluster_0": "term1;term2;term3", ...}
                clusters = data
            else:
                raise ValueError("不支持的JSON格式")
                
        except (json.JSONDecodeError, KeyError, ValueError):
            # 如果不是JSON格式，按原来的文本格式处理
            for line in content.split('\n'):
                if line.strip():  # 跳过空行
                    if '\t' in line:
                        cluster_id, terms = line.strip().split('\t', 1)
                        clusters[cluster_id] = terms
                    else:
                        # 处理其他分隔符，如冒号
                        parts = line.strip().split(':', 1)
                        if len(parts) == 2:
                            clusters[parts[0]] = parts[1]
        
        if not clusters:
            raise ValueError("未能从文件中解析出聚类数据，请检查文件格式")
        
        return clusters
    
    def label_topics(self, cluster_file):
        """
        原有的非流式输出方法，保持向后兼容
        """
        clusters = self._parse_cluster_file(cluster_file)
        
        # 存储标签结果
        labels = []
        cluster_ids = []
        term_lists = []
        
        prompt = "你是一位科技情报信息的分析专家\n你的任务是：\n（1）通过给定的一系列的科技文献的聚类结果及类别下的词，概括每个类的类别。\n（2）根据所有类别的结果，概括这个领域的研究热点是什么。\n输入格式为：cluster_x:主题词1、主题词2、...\n输出格式为：cluster_x的主题标签是《主题标签》\n以下是类别及类别下的词：\n\n"

        # 对每个聚类进行处理
        for cluster_id, terms in clusters.items():
            prompt += f"{cluster_id}:{terms}\n"
            
        label = self.llm_client.get_completion(prompt)
        return label
    
    def label_topics_stream(self, cluster_file):
        """
        流式输出版本的主题标签方法
        返回一个生成器，逐步产生响应内容
        """
        clusters = self._parse_cluster_file(cluster_file)
        
        prompt = "你是一位科技情报信息的分析专家\n你的任务是：\n（1）通过给定的一系列的科技文献的聚类结果及类别下的词，概括每个类的类别。\n（2）根据所有类别的结果，概括这个领域的研究热点是什么。\n输入格式为：cluster_x:主题词1、主题词2、...\n输出格式为：cluster_x的主题标签是《主题标签》\n以下是类别及类别下的词：\n\n"

        # 对每个聚类进行处理
        for cluster_id, terms in clusters.items():
            prompt += f"{cluster_id}:{terms}\n"
            
        # 使用流式接口
        for chunk in self.llm_client.get_completion_stream(prompt):
            yield chunk 